//----------------------------------------------------------------------------
// M A R I T I M E  R E S E A R C H  I N S T I T U T E  N E T H E R L A N D S
//----------------------------------------------------------------------------
// Copyright (C) 2011 - MARIN - All rights reserved - http://www.marin.nl
//----------------------------------------------------------------------------
// Program    : mermaid
// Module     : vbm
// File       : CG.cu
// Author     : M.A. de Jong
//----------------------------------------------------------------------------
//
// CG algorithm with Diagonal scaling and IP preconditioner
//
// Experimental code!! (dec 9, 2011)
//
// TODO code does not work correctly yet, contains bug!
#define USEDIAGSCAL     1
#define USEIPPREC       1
#define USEZEROVECTOR   0
#define USETEXTURES     1
#define TARGET          stdout//m_file
#define CAPTURE_RUNINFO 0

#define NUM_THREADS     256
#define MAX_ITERATIONS  99

#include "../defines.h"

#ifdef USEDOUBLE
  #undef USETEXTURES
  #define USETEXTURES   0
#endif

#include "CG.h"
#include <stdio.h>
//#if CUDART_VERSION < 4000
#include <cuda.h>
//#endif
#include <cublas.h>
#include "../FloatUtils.h"
#include <sys/time.h>
#include <stdio.h>
#include <cstdlib>
#include "ipdiagmisc.h"
#include "kernels/cg_kernels.h"

#define START_CPU_TIMER(time, iter) \
{ \
   cudaThreadSynchronize(); \
	gettimeofday(&start, NULL); \
   for (int t = 0; t < iter; t++) { \

#define STOP__CPU_TIMER(time, iter) \
   } \
   cudaThreadSynchronize(); \
	gettimeofday(&stop,  NULL); \
	time = (stop.tv_sec - start.tv_sec) + \
		1e-6 * (stop.tv_usec - start.tv_usec); \
   time /= iter; \
}

#if USETEXTURES
// texture
texture <float, 1> texRefx;
#endif

__inline__ __device__ REAL fetchX(const REAL *x, int i)
{
#if USETEXTURES
    return tex1Dfetch(texRefx, i);
#else
    return x[i];
#endif
}

// global variables
static REAL m_tolerance;
static REAL *m_dcc;
static REAL *m_dss;
static REAL *m_dww;
#if USEIPPREC
static REAL *m_kcc;
static REAL *m_kss;
static REAL *m_kww;
static REAL *m_dz;
#endif
#if USEDIAGSCAL
static REAL *m_dpp;
#endif
static REAL *m_dr;
static REAL *m_dx;
static REAL *m_dy;
static REAL *m_dp;
static REAL *m_dq;
static int m_numBlocks;
static int m_numThreads;
static size_t m_pitch;
static size_t m_elementSize;
static int m_srcOffset;
static int m_Nx1;
static int m_Nx2;
static int m_n;
#if CAPTURE_RUNINFO
static FILE *m_file;
static int m_counter;
#endif

void setup(Array2D<REAL> *pcc,
           Array2D<REAL> *pcs,
           Array2D<REAL> *pcw,
           int Nx1, int Nx2,
           REAL tolerance)
{
#if CAPTURE_RUNINFO
    m_file = fopen("CUDA_SIMPLE/runinfo_IPDIAG.txt", "w");
    if (m_file == NULL) {
        fprintf(stderr, "Can't create file 'runinfo_IPDIAG.txt'");
        exit(1);
    }

    m_counter = 0;
#endif
 
    m_Nx1 = Nx2;
    m_Nx2 = Nx1;
    m_n = m_Nx1 * m_Nx2;
    m_tolerance = tolerance;

    m_numThreads = NUM_THREADS;

    // No. of threads blocks necessary to handle all data
    m_numBlocks = (int)ceil((float)(m_n) / (float)m_numThreads);

    // number of bytes per array element
    m_elementSize = sizeof(REAL);

    // pitch of the array on the device
    m_pitch = m_Nx1 * m_elementSize;

    // location of first relevant data element
    m_srcOffset = pcc->columns()+1;

    // copy center stencil to device
    m_dcc = ipdiagmisc::initGpuVector<REAL,1>(m_n, 1);
    cudaMemcpy2D(m_dcc, m_pitch, pcc->getElements() + m_srcOffset, 
        pcc->columns() * m_elementSize, m_Nx1 * m_elementSize, 
        m_Nx2, cudaMemcpyHostToDevice);

    // copy south stencil to device
    // Note: on device it becomes the west stencil!
    m_dww = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);
    cudaMemcpy2D(m_dww, m_pitch, pcs->getElements() + m_srcOffset, 
        pcs->columns() * m_elementSize, m_Nx1 * m_elementSize, 
        m_Nx2, cudaMemcpyHostToDevice);

    // copy west stencil to device 
    // Note: on the device it becomes the south stencil!
    m_dss = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);
    cudaMemcpy2D(m_dss, m_pitch, pcw->getElements() + m_srcOffset, 
        pcw->columns() * m_elementSize, m_Nx1 * m_elementSize, 
        m_Nx2, cudaMemcpyHostToDevice);

    // creating a CUDA stream
    // Note: not necessary in case kernels do not overlap
    cudaStream_t stream;
    cudaStreamCreate(&stream);


#if USEDIAGSCAL
    // allocate memory for the diagonal matrix on device
    m_dpp = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

    // make diagonal matrix, to be precise: sqrt(D), where D = diag(A)
    cg::makeDiagonal<REAL> <<< m_numBlocks,        // No. thread blocks
                               m_numThreads,       // No. threads per block
                               0,                  // Amount of shared memory
                               stream              // Use this CUDA stream
                           >>> (m_dpp, m_dcc, m_n);

    // update matrix A
    cg::updateMatrix<REAL> <<< m_numBlocks, m_numThreads, 0, stream >>>  
                               (m_dcc, m_dss, m_dww, m_dpp, m_Nx1, m_Nx2);
#endif

#if USEIPPREC
    m_kcc = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);
    m_kss = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);
    m_kww = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

    // make preconditioning matrix K from (possibly updated) matrix A
    cg::makeIPPrec<REAL> <<< m_numBlocks, m_numThreads, 0, stream >>> 
                         (m_kcc, m_kss, m_kww, m_dcc, m_dss, m_dww, Nx1, Nx2);

    // allocate vector z on device
    m_dz = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);
#endif

    // allocate vector y on device
    m_dy = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

    // allocate vector p on device
    m_dp = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

    // allocate vector q on device
    m_dq = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

    // allocate vector r on device
    m_dr = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

    // allocate vector x on device
    m_dx = ipdiagmisc::initGpuVector<REAL,1>(m_n, 0);

#if CAPTURE_RUNINFO
    fprintf(TARGET, "CG output per frame:\n");
    fprintf(TARGET, "\n");
    fprintf(TARGET, "Frame:\t|<r,z>|\t\t||r||\t\t||b||"
                    "\t      iter\t  time\n");
    fprintf(TARGET, "=====================================");
    fprintf(TARGET, "=====================================\n");
#endif
}

void start(Array2D<REAL> *pB,
           Array2D<REAL> *pX)
{
#if CAPTURE_RUNINFO
    ++m_counter;
    struct timeval start, stop;
    float timeCGiter;
    START_CPU_TIMER(time, 1);
#endif

    int iteration;
    int maxIterations = MAX_ITERATIONS;
    REAL rhonew;
    REAL rhoold;
    REAL rhobew;
    REAL beta;
    REAL alpha;
    REAL sigma;

    // creating a CUDA stream
    // Note: not strictly necessary when kernels do not overlap
    cudaStream_t stream1;
    cudaStream_t stream2;
    cudaStreamCreate(&stream1);
    cudaStreamCreate(&stream2);

    // copy vector b to device (to r)
    cudaMemcpy2DAsync(m_dr, m_pitch, pB->getElements() + m_srcOffset, 
        pB->columns() * m_elementSize, m_Nx1 * m_elementSize, 
        m_Nx2, cudaMemcpyHostToDevice, stream1);

#if CAPTURE_RUNINFO
    // ||b||
  #ifdef USEDOUBLE
    REAL normRHS = cublasDdot(m_n, m_dr, 1, m_dr, 1);
  #else
    REAL normRHS = cublasSdot(m_n, m_dr, 1, m_dr, 1);
  #endif
#endif

#if USEZEROVECTOR
    ipdiagmisc::memset_kernel<REAL> <<< m_numBlocks, m_numThreads, 0, stream2 >>> 
                                  (m_dx, m_n, 0);
#endif

#if USEDIAGSCAL
    // x = P^Tx
    cg::elementWiseMul<REAL>  <<< m_numBlocks, m_numThreads, 0, stream2 >>> 
                                  (m_dx, m_dpp, m_n);

    // y = A * x
  #if USETEXTURES
    cudaBindTexture(NULL, texRefx, m_dx);
  #endif
    cg::SpMVv2<REAL> <<< m_numBlocks, m_numThreads, 0, stream2 >>> 
                         (m_dy, m_dcc, m_dss, m_dww, m_dx, m_Nx1, m_Nx2);
#else
  #if USETEXTURES
    cudaBindTexture(NULL, texRefx, m_dx);
  #endif
    cg::SpMVv1<REAL> <<< m_numBlocks, m_numThreads, 0, stream2 >>> 
                         (m_dy, m_dcc, m_dss, m_dww, m_dx, m_Nx1, m_Nx2);
#endif

#if USEDIAGSCAL
    // r = P^{-1} * r = P^{-T} * r
    cg::elementWiseDiv<REAL> <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                                 (m_dr, m_dpp, m_n);
#endif

    // r = r - y
    cg::AXPY<REAL>   <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                         (m_dr, m_dy, -1, 1, m_n);

#if CAPTURE_RUNINFO
    // ||r|| = ||b - A*x||
  #ifdef USEDOUBLE
    REAL normRes = cublasDdot(m_n, m_dr, 1, m_dr, 1);
  #else
    REAL normRes = cublasSdot(m_n, m_dr, 1, m_dr, 1);
  #endif
#endif

#if USEIPPREC
    // z = M^{-1} * r
  #if USETEXTURES
    cudaBindTexture(NULL, texRefx, m_dr);
  #endif
    cg::SpMVv1<REAL> <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                         (m_dz, m_kcc, m_kss, m_kww, m_dr, m_Nx1, m_Nx2);
#endif

    //START_CPU_TIMER(time, 1);

#if USEIPPREC
  #ifdef USEDOUBLE
    rhonew = cublasDdot(m_n, m_dr, 1, m_dz, 1);
  #else
    rhonew = cublasSdot(m_n, m_dr, 1, m_dz, 1);
  #endif
#else
  #ifdef USEDOUBLE
    rhonew = cublasDdot(m_n, m_dr, 1, m_dr, 1);
  #else
    rhonew = cublasSdot(m_n, m_dr, 1, m_dr, 1);
  #endif
#endif

#if CAPTURE_RUNINFO
    fprintf(TARGET, "%d\t%1.6E\t%1.6E\t%1.6E", 
        m_counter, rhonew, normRes, normRHS);
#endif

    // NOTE: because here all nodes count, the stopping criterium is slightly
    // different from the one in the RRB-method (which operates on red nodes)
    rhobew = (rhonew + 1.0) * FloatUtils::square(m_tolerance); 

    iteration = 0;

    // here the CG iterations start:
    while (fabs(rhonew) > fabs(rhobew) && iteration < maxIterations) 
    {
        ++iteration;

        // for first iteration a different treatment:
        if (iteration == 1) 
        {
#if USEIPPREC
            // p = z
            ipdiagmisc::memcpy<REAL> (m_dp, m_dz, m_n);
#else
            // p = r
            ipdiagmisc::memcpy<REAL> (m_dp, m_dr, m_n);
#endif
        }
        else // for all iterations apart from the first one:
        {
            // beta = rhonew / rhoold
            beta = rhonew / rhoold;
#if USEIPPREC
            // p = z + beta * p
            cg::AXPY<REAL> <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                               (m_dp, m_dz, 1, beta, m_n);
#else
            // p = r + beta * p
            cg::AXPY<REAL> <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                               (m_dp, m_dr, 1, beta, m_n);
#endif
        }

        // q = A * p
#if USETEXTURES
        cudaBindTexture(NULL, texRefx, m_dp);
#endif
#if USEDIAGSCAL
        cg::SpMVv2<REAL>   <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                               (m_dq, m_dcc, m_dss, m_dww, m_dp, m_Nx1, m_Nx2);
#else
        cg::SpMVv1<REAL>   <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                               (m_dq, m_dcc, m_dss, m_dww, m_dp, m_Nx1, m_Nx2);
#endif

        // sigma = <p,q>
#ifdef USEDOUBLE
        sigma = cublasDdot(m_n, m_dp, 1, m_dq, 1);
#else
        sigma = cublasSdot(m_n, m_dp, 1, m_dq, 1);
#endif

        // alpha = rhonew / sigma
        alpha = rhonew / sigma;

        // r = r - alpha * q
        cg::AXPY<REAL>     <<< m_numBlocks, m_numThreads, 0, stream2 >>> 
                               (m_dr, m_dq, -alpha, 1, m_n);

        // x = x + alpha * p
        cg::AXPY<REAL>     <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                               (m_dx, m_dp,  alpha, 1, m_n);

#if USEIPPREC
        // z = M^{-1} * r
  #if USETEXTURES
        cudaBindTexture(NULL, texRefx, m_dr);
  #endif
        cg::SpMVv1<REAL>   <<< m_numBlocks, m_numThreads, 0, stream2 >>> 
                               (m_dz, m_kcc, m_kss, m_kww, m_dr, m_Nx1, m_Nx2);
#endif
        // rhoold = rhonew
        rhoold = rhonew;

#if USEIPPREC
        // rhonew = <r,z>
  #ifdef USEDOUBLE
        rhonew = cublasDdot(m_n, m_dr, 1, m_dz, 1);
  #else
        rhonew = cublasSdot(m_n, m_dr, 1, m_dz, 1);
  #endif
#else
        // rhonew = <r,r>
  #ifdef USEDOUBLE
        rhonew = cublasDdot(m_n, m_dr, 1, m_dr, 1);
  #else
        rhonew = cublasSdot(m_n, m_dr, 1, m_dr, 1);
  #endif
#endif
    }
#if CAPTURE_RUNINFO
    fprintf(TARGET, "\t%d", iteration);
#endif


#if USEDIAGSCAL
    // x = P^{-T} * x
    cg::elementWiseDiv<REAL> <<< m_numBlocks, m_numThreads, 0, stream1 >>> 
                             (m_dx, m_dpp, m_n);
#endif

    // copy vector x to host
    cudaMemcpy2D(pX->getElements() + m_srcOffset, 
        pX->columns() * m_elementSize, m_dx, m_Nx1 * m_elementSize, 
        m_Nx1 * m_elementSize, m_Nx2, cudaMemcpyDeviceToHost);

#if CAPTURE_RUNINFO
    STOP__CPU_TIMER(timeCGiter, 1);
    fprintf(TARGET, "\t%3.6f\n", timeCGiter);
#endif
}



void clean()
{
#if CAPTURE_RUNINFO
    // close file
    fclose(m_file);
#endif
    cudaFree(m_dcc);
    cudaFree(m_dss);
    cudaFree(m_dww);
#if USEIPPREC
    cudaFree(m_kcc);
    cudaFree(m_kss);
    cudaFree(m_kww);
    cudaFree(m_dz);
#endif
#if USEDIAGSCAL
    cudaFree(m_dpp);
#endif
    cudaFree(m_dr);
    cudaFree(m_dx);
    cudaFree(m_dy);
    cudaFree(m_dp);
    cudaFree(m_dq);
}

